#ifndef IO_CONTEXT_POOL_H
#define IO_CONTEXT_POOL_H

#include <asio.hpp>
#include <list>
#include <memory>
#include <vector>
#include <cstdint>

class IoContextPool
{
public:
    static IoContextPool &Instance(uint32_t io_context_num)
    {
        static IoContextPool singleton(io_context_num);
        return singleton;
    }

    /// Get an io_context to use.
    asio::io_context &GetIoContext()
    {
        asio::io_context &io_context = *io_contexts_[(++next_io_context_) % io_context_num_];
        return io_context;
    }

    IoContextPool(const IoContextPool &) = delete;
    IoContextPool &operator=(const IoContextPool &) = delete;

private:
    typedef asio::executor_work_guard<asio::io_context::executor_type> io_context_work;

    explicit IoContextPool(uint32_t io_context_num)
    {
        if (0 == io_context_num)
        {
            io_context_num_ = 1;
        }
        else
        {
            io_context_num_ = io_context_num;
        }

        // Give all the io_contexts work to do so that their run() functions will not
        // exit until they are explicitly stopped.
        for (uint32_t i = 0; i < io_context_num_; ++i)
        {
            const std::shared_ptr<asio::io_context> io_context(new asio::io_context);
            io_contexts_.push_back(io_context);
            work_.push_back(asio::make_work_guard(*io_context));
        }

        run();
    }

    ~IoContextPool()
    {
        stop();
    }

    /// Stop all io_context objects in the pool.
    void stop()
    {
        // Explicitly stop all io_contexts.
        for (uint32_t i = 0; i < io_context_num_; ++i)
        {
            io_contexts_[i]->stop();
        }

        // Wait for all threads in the pool to exit.
        for (uint32_t i = 0; i < threads_.size(); ++i)
        {
            if (threads_[i].joinable())
            {
                threads_[i].join();
            }
        }
    }

    /// Run all io_context objects in the pool.
    void run()
    {
        // Create a pool of threads to run all of the io_contexts.
        for (uint32_t i = 0; i < io_context_num_; ++i)
        {
            for (uint32_t k = 0; k < NUMBER_OF_THREADS_PER_CONTEXT; ++k)
            {
                threads_.emplace_back([this, i] {
                    io_contexts_[i]->run();
                });
            }
        }
    }

    uint32_t io_context_num_ = 0;
    const uint32_t NUMBER_OF_THREADS_PER_CONTEXT = 3;
    std::vector<std::thread> threads_;
    std::vector<std::shared_ptr<asio::io_context>> io_contexts_;
    std::list<io_context_work> work_;
    std::atomic<uint32_t> next_io_context_ = 0;
};

#endif